#!/usr/bin/env python
# coding: utf-8
# 连续遗传算法，寻找Ackley函数的最小值

import numpy as np
import matplotlib.pyplot as plt
import itertools
import random
import math


class Individual:
    def __init__(self, genetype=None, fitness=None):
        self.genetype = genetype
        self.fitness = fitness

    def __str__(self):
        return '{0},{1}'.format(self.genetype, self.fitness)


# 二维Ackley函数
def ackley(X):
    x = X[0]
    y = X[0]
    return math.e - 20 * math.exp(-0.2 * math.sqrt((x * x + y * y) / 2)) - math.exp(
        (math.cos(2 * math.pi * x) + math.cos(2 * math.pi * y)) / 2)


# 20维sphere函数
def sphere(X):
    return np.sum(X * X)


def fitness(pop):
    for indv in pop:
        # value = -ackley(*indv.genetype)
        value = 1.0 / sphere(indv.genetype)
        # value = -sphere(indv.genetype)
        # assert value >= 0
        indv.fitness = value


# 轮盘赌算法
def roulette(values, _sum):
    rand = random.random() * _sum
    f = values[0]
    k = 0
    while f < rand:
        k += 1
        f += values[k]
    return k


# 选出一对父代
def selection(pop):
    # 处理负数，临时存储，不改变原来的数值
    fitness_values = np.array([indv.fitness for indv in pop])
    fitness_values -= np.min(fitness_values)
    # 求和
    _sum = np.sum(fitness_values)
    j = roulette(fitness_values, _sum)
    k = roulette(fitness_values, _sum)
    # 不能自己交配
    while k == j:
        k = roulette(fitness_values, _sum)

    return pop[j], pop[k]


def cross(parents):
    # 保证有交叉，限制范围
    rand = random.randrange(1, len(parents[0].genetype))
    # print(parents[0].genetype, parents[1].genetype)
    child1_gene = np.append(parents[0].genetype[:rand], parents[1].genetype[rand:])
    child2_gene = np.append(parents[1].genetype[:rand], parents[0].genetype[rand:])
    # print(child1_gene, child2_gene)
    child1 = Individual(child1_gene)
    child2 = Individual(child2_gene)
    return child1, child2


def mutate(children_pop):
    for child in children_pop:
        new_genetype = []
        for i, bit in enumerate(child.genetype):
            rand = random.random()
            if rand < mutate_rate:  # mutate
                new_genetype.append(random.uniform(xi_min, xi_max))
            else:
                new_genetype.append(bit)
        child.genetype = np.array(new_genetype)


# def visualize(pop, total_count, count):
#     row = total_count / 2
#     plt.subplot(row, 2, count)
#     plt.xlim(-5, 5)
#     plt.ylim(-5, 5)
#     points = [p.genetype for p in pop]
#     plt.scatter(*zip(*points))
#

pop_size = 20
mutate_rate = 0.05
xi_min = -5
xi_max = 5
dim = 20
# initialize
# parents_pop = [Individual([random.uniform(x_min, x_max), random.uniform(y_min, y_max)]) for i in range(pop_size)]
parents_pop = []
for i in range(pop_size):
    genetype = []
    for j in range(dim):
        genetype.append(random.uniform(xi_min, xi_max))
    parents_pop.append(Individual(np.array(genetype)))
print(parents_pop[0])

generation = 100
count = 0

while count < generation:
    # 计算种群中每个父代的适应度
    fitness(parents_pop)
    # visualize(parents_pop, generation, count + 1)
    fitness_values = [p.fitness for p in parents_pop]
    print('gen: {0}, best: {1}, avg: {2}, sphere: {3}'.format(count, max(fitness_values), np.mean(fitness_values),
                                                              1.0 / np.mean(fitness_values)))
    children_pop = []
    while len(children_pop) < len(parents_pop):
        # 用适应度根据概率选出一对交配的父代
        parents = selection(parents_pop)
        # 交配生成子代
        children = cross(parents)
        children_pop += children
    # 变异
    mutate(children_pop)

    parents_pop = children_pop
    count += 1

plt.show()
